# tushare_mirror/storage.py
import os
import pandas as pd
import pymysql
from typing import List, Dict, Any, Optional
from vnpy.trader.object import BarData
from vnpy.trader.constant import Exchange, Interval
from .config import Config
from .exceptions import StorageError


class DataStorage:
    def __init__(self):
        self.data_path = Config.DATA_PATH
        self.db_config = {
            "host": Config.DB_HOST,
            "port": Config.DB_PORT,
            "user": Config.DB_USER,
            "password": Config.DB_PASSWORD,
            "database": Config.DB_NAME,
            "charset": "utf8mb4",
            "cursorclass": pymysql.cursors.DictCursor
        }
        os.makedirs(self.data_path, exist_ok=True)
        self._init_db_tables()

    def _init_db_tables(self) -> None:
        with pymysql.connect(**self.db_config) as conn:
            with conn.cursor() as cursor:
                # 修正建表语句
                cursor.execute("""
                                CREATE TABLE IF NOT EXISTS bars (
                                    id INT AUTO_INCREMENT PRIMARY KEY,
                                    symbol VARCHAR(10),
                                    exchange VARCHAR(10),
                                    datetime DATETIME,
                                    `interval` VARCHAR(5),
                                    open DECIMAL(18,4),
                                    high DECIMAL(18,4),
                                    low DECIMAL(18,4),
                                    close DECIMAL(18,4),
                                    volume DECIMAL(18,4),
                                    INDEX idx_symbol (symbol),
                                    INDEX idx_datetime (datetime)
                                )
                            """)

                cursor.execute("""
                    CREATE TABLE IF NOT EXISTS valuation (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        symbol VARCHAR(10),
                        pe DECIMAL(18,4),
                        pb DECIMAL(18,4),
                        ps DECIMAL(18,4),
                        dv_ratio DECIMAL(18,4),
                        circ_mv DECIMAL(18,4),
                        `datetime` DATETIME,  # 转义保留字
                        INDEX idx_symbol (symbol),
                        INDEX idx_datetime (`datetime`)  # 转义索引字段
                    )
                """)
                conn.commit()

    def save_bars(self, bars: List[BarData], symbol: str) -> None:
        """保存日线数据（CSV+MySQL）"""
        try:
            self._save_bars_to_csv(bars, symbol)
            self._save_bars_to_mysql(bars)
        except Exception as e:
            raise StorageError(f"日线保存失败: {str(e)}") from e

    def _save_bars_to_csv(self, bars: List[BarData], symbol: str) -> None:
        """CSV保存实现"""
        data = [{
            "symbol": bar.symbol,
            "exchange": bar.exchange.value,
            "datetime": bar.datetime.strftime("%Y-%m-%d %H:%M:%S"),
            "interval": bar.interval.value,
            "open_price": bar.open_price,
            "high_price": bar.high_price,
            "low_price": bar.low_price,
            "close_price": bar.close_price,
            "volume": bar.volume
        } for bar in bars]

        df = pd.DataFrame(data)
        path = os.path.join(self.data_path, f"{symbol}_daily.csv")

        # 处理首次写入列名，后续追加
        header = not os.path.exists(path)
        df.to_csv(path, mode='a', index=False, header=header)

    def _save_bars_to_mysql(self, bars: List[BarData]) -> None:
        with pymysql.connect(**self.db_config) as conn:
            with conn.cursor() as cursor:
                # 修正插入语句
                sql = """INSERT INTO bars 
                    (symbol, exchange, datetime, `interval`, 
                     open, high, low, close, volume)
                    VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)"""

                params = [
                    (bar.symbol, bar.exchange.value, bar.datetime,
                     bar.interval.value, bar.open_price, bar.high_price,
                     bar.low_price, bar.close_price, bar.volume)
                    for bar in bars
                ]

                cursor.executemany(sql, params)
                conn.commit()

    def save_valuation(self, df: pd.DataFrame, symbol: str) -> None:
        """保存估值数据（CSV+MySQL）"""
        try:
            self._save_valuation_to_csv(df, symbol)
            self._save_valuation_to_mysql(df)
        except Exception as e:
            raise StorageError(f"估值保存失败: {str(e)}") from e

    def _save_valuation_to_csv(self, df: pd.DataFrame, symbol: str) -> None:
        path = os.path.join(self.data_path, f"{symbol}_valuation.csv")
        header = not os.path.exists(path)
        df.to_csv(path, mode='a', index=False, header=header)

    def _save_valuation_to_mysql(self, df: pd.DataFrame) -> None:
        with pymysql.connect(**self.db_config) as conn:
            with conn.cursor() as cursor:
                sql = """INSERT INTO valuation 
                    (symbol, pe, pb, ps, dv_ratio, circ_mv, datetime)
                    VALUES (%s, %s, %s, %s, %s, %s, %s)"""

                # 逐行处理 NaN 值
                params = []
                for _, row in df.iterrows():
                    param = (
                        row['symbol'],
                        None if pd.isna(row['pe']) else round(row['pe'], 4),
                        None if pd.isna(row['pb']) else round(row['pb'], 4),
                        None if pd.isna(row['ps']) else round(row['ps'], 4),
                        None if pd.isna(row['dv_ratio']) else round(row['dv_ratio'], 4),
                        None if pd.isna(row['circ_mv']) else round(row['circ_mv'], 4),
                        row['datetime']
                    )
                    params.append(param)

                cursor.executemany(sql, params)
                conn.commit()

    def get_latest_bar(self, symbol: str) -> Optional[BarData]:
        """从数据库获取最新日线数据"""
        with pymysql.connect(**self.db_config) as conn:
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT * FROM bars 
                    WHERE symbol = %s 
                    ORDER BY datetime DESC 
                    LIMIT 1
                """, (symbol,))
                result = cursor.fetchone()
                return self._convert_db_row_to_bar(result) if result else None

    def _convert_db_row_to_bar(self, row: dict) -> BarData:
        """数据库行转BarData对象"""
        return BarData(
            symbol=row["symbol"],
            exchange=Exchange(row["exchange"]),
            datetime=row["datetime"],
            interval=Interval(row["interval"]),
            open_price=row["open"],
            high_price=row["high"],
            low_price=row["low"],
            close_price=row["close"],
            volume=row["volume"],
            gateway_name="tushare"
        )

    def get_latest_valuation(self, symbol: str) -> Optional[dict]:
        """从数据库获取最新估值数据"""
        with pymysql.connect(**self.db_config) as conn:
            with conn.cursor() as cursor:
                cursor.execute("""
                    SELECT * FROM valuation 
                    WHERE symbol = %s 
                    ORDER BY datetime DESC 
                    LIMIT 1
                """, (symbol,))
                return cursor.fetchone()

    def has_data(self, symbol: str, data_type: str) -> bool:
        """检查是否存在指定数据"""
        if data_type == "daily":
            csv_path = os.path.join(self.data_path, f"{symbol}_daily.csv")
            return os.path.exists(csv_path)
        # 可以扩展其他数据类型检查
        return False